package com.xiashitech.agent.instrumentation.spring.security.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.dto.UserRequest;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSSecurityContextPersistenceCallDepth;
import com.xiashitech.agent.instrumentation.spring.security.context.SecurityRequestContext;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import javax.servlet.ServletRequest;
import javax.servlet.http.HttpServletRequest;
import java.util.Date;

import static com.xiashitech.agent.instrumentation.spring.security.XSSpringSecuritySingletons.instrumenter;
import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSSecurityContextPersistenceInstrumentation implements TypeInstrumentation {
    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return named("org.springframework.security.web.context.SecurityContextPersistenceFilter");
    }

    @Override
    public void transform(TypeTransformer transformer) {
        transformer.applyAdviceToMethod(
            isMethod()
                .and(isPublic())
                .and(named("doFilter"))
                .and(takesArgument(0, named("javax.servlet.ServletRequest")))
                .and(takesArgument(1, named("javax.servlet.ServletResponse")))
                .and(takesArgument(2, named("javax.servlet.FilterChain")))
                .and(takesArguments(3)),
                XSSecurityContextPersistenceInstrumentation.class.getName() + "$PersistenceAdvice");
    }
    public static class PersistenceAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.FieldValue("FILTER_APPLIED") String filterApplied,
                @Advice.Argument(0) ServletRequest request,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope) {
            DurationPrinter durationPrinter = new DurationPrinter("XSSecurityContextPersistenceInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSSecurityContextPersistenceCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            if(request instanceof HttpServletRequest && request.getAttribute(filterApplied) == null) {
                SecurityRequestContext.setSecurityUserRequest(UserRequest.builder().requestDate(String.valueOf(new Date().getTime())).build());
                try {
                    Context parentContext = currentContext();
                    if (instrumenter().shouldStart(parentContext,null)) {
                        parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                        context = instrumenter().start(parentContext, null);
                        context = BaggageUtil.syncBaggage().build().storeInContext(context);
                        scope = context.makeCurrent();
                    }
                } catch (Throwable e) {
                    StartConfig.throwableStackPrinter(e);
                }
            }
            durationPrinter.printElapseDuration("PersistenceAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope) {
            DurationPrinter durationPrinter = new DurationPrinter("XSSecurityContextPersistenceInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if (callDepth.decrementAndGet() > 0) {
                return;
            }
            if (scope == null) {
                return;
            }
            scope.close();
            UserRequest securityUserRequest = SecurityRequestContext.getSecurityUserRequest();
            instrumenter().end(context,null,securityUserRequest,null);
            SecurityRequestContext.refresh();
            durationPrinter.printElapseDuration("PersistenceAdvice e");
        }
    }
}
